package handler

import (
	"github.com/miekg/dns"
	"lktx-dns/config"
	"lktx-dns/entity"
	"lktx-dns/stat"
	"log"
	"math/rand"
	"net"
	"strings"
	"time"
)

type LktxDNSHandler struct {
}

func (receiver LktxDNSHandler) removeLastDot(s string) string {
	// 检查最后一个字符是否是 "."
	if strings.HasSuffix(s, ".") {
		// 删除最后一个字符
		return s[:len(s)-1]
	}
	return s
}

func (h *LktxDNSHandler) ServeDNS(w dns.ResponseWriter, r *dns.Msg) {
	msg := new(dns.Msg)
	msg.SetReply(r)
	domain := r.Question[0].Name
	//log.Printf("解析:%s", domain)
	// 查找配置文件中是否有该域名的记录
	value, ok := config.Records.Load(h.removeLastDot(domain))
	if !ok {
		stat.StatData.AddDNSError(domain)
		res, err := dns.Exchange(r, config.ConfigData.Dns.Default)
		if err != nil {
			//log.Printf("使用默认解析错误 %s %s: %v", config.ConfigData.Dns.Default, domain, err)
			err := w.WriteMsg(msg)
			if err != nil {
				log.Printf(err.Error())
				return
			}
			return
		}
		msg.Answer = append(msg.Answer, res.Answer...)
		err = w.WriteMsg(msg)
		if err != nil {
			log.Printf(err.Error())
			return
		}
		return
	}

	entities := value.([]*entity.DomainInfoEntity)
	// 根据请求的类型返回相应的记录
	switch r.Question[0].Qtype {
	case dns.TypeTXT:
		var texts []string
		for _, infoEntity := range entities {
			if *infoEntity.Type == `TXT` {
				texts = append(texts, *infoEntity.Record)
			}
		}
		if len(texts) != 0 {
			txtRecord := &dns.TXT{
				Hdr: dns.RR_Header{
					Name:   domain,
					Rrtype: dns.TypeTXT,
					Class:  dns.ClassINET,
					Ttl:    300,
				},
				Txt: texts,
			}
			msg.Answer = append(msg.Answer, txtRecord)
		}

	case dns.TypeA:
		for _, infoEntity := range entities {
			if *infoEntity.Type == `A` {
				rr := &dns.A{
					Hdr: dns.RR_Header{
						Name:   domain,
						Rrtype: dns.TypeA,
						Class:  dns.ClassINET,
						Ttl:    300,
					},
					A: net.ParseIP(*infoEntity.Record),
				}
				msg.Answer = append(msg.Answer, rr)
			}
		}

	case dns.TypeMX:
		for _, infoEntity := range entities {
			if *infoEntity.Type == `MX` {
				mxRecord := &dns.MX{
					Hdr: dns.RR_Header{
						Name:   domain,
						Rrtype: dns.TypeMX,
						Class:  dns.ClassINET,
						Ttl:    300,
					},
					Mx:         *infoEntity.Record + ".",
					Preference: uint16(10),
				}
				msg.Answer = append(msg.Answer, mxRecord)
			}
		}

	case dns.TypeCNAME:
		for _, infoEntity := range entities {
			if *infoEntity.Type == `CNAME` {
				cnameRecord := &dns.CNAME{
					Hdr: dns.RR_Header{
						Name:   domain,
						Rrtype: dns.TypeCNAME,
						Class:  dns.ClassINET,
						Ttl:    300,
					},
					Target: *infoEntity.Record + ".",
				}
				msg.Answer = append(msg.Answer, cnameRecord)
			}
		}
	}

	if len(msg.Answer) == 0 {
		stat.StatData.AddDNSError(domain)
		res, err := dns.Exchange(r, config.ConfigData.Dns.Default)
		if err != nil {
			//log.Printf("使用默认解析错误 %s %s: %v", config.ConfigData.Dns.Default, domain, err)
			err := w.WriteMsg(msg)
			if err != nil {
				log.Printf(err.Error())
				return
			}
			return
		}
		msg.Answer = append(msg.Answer, res.Answer...)
	}
	stat.StatData.AddDNSSuccess(domain)
	rand.New(rand.NewSource(time.Now().UnixNano()))
	// 使用 Shuffle 打乱切片的顺序
	rand.Shuffle(len(msg.Answer), func(i, j int) {
		msg.Answer[i], msg.Answer[j] = msg.Answer[j], msg.Answer[i]
	})

	err := w.WriteMsg(msg)
	if err != nil {
		log.Printf(err.Error())
		return
	}
}
